import os.path

import torch

from data.base_dataset import BaseDataset, get_params, get_transform
from data.image_folder import make_dataset
from PIL import Image
import torchvision.transforms as transforms


class AlignedDataset(BaseDataset):
    def initialize(self, opt):
        self.opt = opt
        self.root = opt.dataroot
        self.phase = opt.phase
        self.isVal = False

        if opt.phase == 'val':
            self.phase = 'test'
            self.isVal = True

        # input A (label maps)
        dir_A = '/SAR'
        self.dir_A = os.path.join(opt.dataroot, self.phase + dir_A)
        self.A_paths = sorted(make_dataset(self.dir_A))

        # input B (real images)
        if opt.isTrain or self.isVal or opt.use_encoded_image:
            dir_B = '/RGB'
            self.dir_B = os.path.join(opt.dataroot, self.phase + dir_B)
            self.B_paths = sorted(make_dataset(self.dir_B))

        self.dataset_size = len(self.A_paths)

    def __getitem__(self, index):
        # input A (label maps)
        A_path = self.A_paths[index]
        A = Image.open(A_path)
        params = get_params(self.opt, A.size)

        transform_A = get_transform(self.opt, params, normalize=False)
        A_tensor = transform_A(A.convert('L'))

        # 归一化
        mono_transform = transforms.Normalize(mean=[0.5], std=[0.5])
        A_tensor = mono_transform(A_tensor)

        B_tensor = inst_tensor = feat_tensor = torch.empty(0)
        if self.opt.isTrain or self.isVal or self.opt.use_encoded_image:
            B_path = self.B_paths[index]
            B = Image.open(B_path).convert('RGB')
            transform_B = get_transform(self.opt, params)
            B_tensor = transform_B(B)

        input_dict = {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor,
                      'feat': feat_tensor, 'path': A_path}

        return input_dict

    def __len__(self):
        return len(self.A_paths) // self.opt.batchSize * self.opt.batchSize

    def name(self):
        return 'AlignedDataset'
